#pragma once

#include <memory>
#include <vector>
#include <string>
#include <cstdint>
#include <cstddef>
#include <ranges>
#include <zenopy/UsdAPI.hpp>
#include <zenox/glm.hpp>

namespace zenopy {

struct UsdPrimData {
    virtual ~UsdPrimData() = default;
};

struct UsdPrimDataXform : UsdPrimData {
    glm::mat4x4 xformMatrix;
};

struct UsdPrimDataMesh : UsdPrimData {
    std::vector<uint32_t> faceVertexCounts;
    std::vector<uint32_t> faceVertexIndices;
    std::vector<glm::vec3> points;
};

#define ZENO_FOREACH_USD_PRIM_DATA_TYPE(_F) \
    /* X-macro pattern begin */ \
    _F(Xform) \
    _F(Mesh) \
    /* X-macro pattern end */

struct UsdPrim {
    using Ptr = std::unique_ptr<UsdPrim>;

    static Ptr make() {
        return std::make_unique<UsdPrim>();
    }

    std::string path;
    std::vector<UsdPrim::Ptr> children;
    std::unique_ptr<UsdPrimData> dataPtr;
    std::string dataType;

    void setDataType(std::string const &newDataType) {
        dataType = newDataType;
#define _FUNCTION(_type) \
        if (dataType == #_type) { \
            dataPtr = std::make_unique<UsdPrimData##_type>(); \
        } else
        ZENO_FOREACH_USD_PRIM_DATA_TYPE(_FUNCTION)
#undef _FUNCTION
        {
            zenox::log_warn("unrecognized data type: {}", dataType);
        }
    }

    template <class UsdPrimDataDerived>
    UsdPrimDataDerived *unsafeCastDataPtr() const {
        static_assert(std::is_base_of_v<UsdPrimData, UsdPrimDataDerived>);
        return static_cast<UsdPrimDataDerived *>(dataPtr.get());
    }

    template <class UsdPrimDataDerived>
    UsdPrimDataDerived *tryCastDataPtr() const {
        static_assert(std::is_base_of_v<UsdPrimData, UsdPrimDataDerived>);
        return dynamic_cast<UsdPrimDataDerived *>(dataPtr.get());
    }
};

struct UsdStage {
    UsdPrim::Ptr rootPrim;
};

namespace details {

inline glm::mat4x4 xformOpMatrix(glm::mat4x4 &m, PythonClient &cli, std::string const &stageFile, std::string const &currPrimPath, std::string_view op) {
    bool inverted = op.starts_with("!invert!");
    if (inverted) {
        op.remove_prefix(8);
    }
    auto strOp = std::string(op);
#define PER_XFORM_OP(_name, _type, _matgen) \
    if (op == #_name) { \
        auto value = stagePrimPropValScalar<_type>(cli, stageFile, currPrimPath, strOp); \
        m = _matgen; \
    } else
    PER_XFORM_OP(translate, glm::dvec3, glm::translate(m, glm::vec3(!inverted ? value : -value)))
    PER_XFORM_OP(translate:pivot, glm::vec3, glm::translate(m, !inverted ? value : -value))
    PER_XFORM_OP(rotateXYZ, glm::vec3, !inverted ? glm::rotate(glm::rotate(glm::rotate(m, value.x, glm::vec3(1, 0, 0)), value.y, glm::vec3(0, 1, 0)), value.z, glm::vec3(0, 0, 1)) : glm::rotate(glm::rotate(glm::rotate(m, -value.z, glm::vec3(0, 0, 1)), -value.y, glm::vec3(0, 1, 0)), -value.x, glm::vec3(1, 0, 0)))
    PER_XFORM_OP(rotateX, float, glm::rotate(m, inverted ? value : -value, glm::vec3(1, 0, 0)))
    PER_XFORM_OP(rotateY, float, glm::rotate(m, inverted ? value : -value, glm::vec3(0, 1, 0)))
    PER_XFORM_OP(rotateZ, float, glm::rotate(m, inverted ? value : -value, glm::vec3(0, 0, 1)))
    PER_XFORM_OP(scale, glm::vec3, glm::scale(m, inverted ? value : -value))
    {
    }
    return {};
}

inline UsdPrim::Ptr recursivelyLoadPrimAndItsChildren(PythonClient &cli, std::string const &stageFile, std::string const &currPrimPath) {
    auto currPrimObj = UsdPrim::make();
    {
        currPrimObj->path = currPrimPath;
        auto dataType = stagePrimType(cli, stageFile, currPrimPath);
        currPrimObj->setDataType(dataType);
        if (dataType == "Mesh") {
            auto mesh = currPrimObj->unsafeCastDataPtr<UsdPrimDataMesh>();
            stagePrimPropValArray(cli, stageFile, currPrimPath, "faceVertexCounts", mesh->faceVertexCounts);
            stagePrimPropValArray(cli, stageFile, currPrimPath, "faceVertexIndices", mesh->faceVertexIndices);
            stagePrimPropValArray(cli, stageFile, currPrimPath, "points", mesh->points);
        } else if (dataType == "Xform") {
            auto xform = currPrimObj->unsafeCastDataPtr<UsdPrimDataXform>();
            auto ops = stagePrimPropValStrList(cli, stageFile, currPrimPath, "xformOpOrder");
            glm::mat4x4 m(1);
            for (std::string_view op: ops | std::views::reverse) {
                xformOpMatrix(m, cli, stageFile, currPrimPath, op);
            }
            xform->xformMatrix = m;
        }
    }
    for (auto const &chPrimPath: stagePrimChildren(cli, stageFile, currPrimPath)) {
        auto chPrimObj = recursivelyLoadPrimAndItsChildren(cli, stageFile, chPrimPath);
        currPrimObj->children.push_back(std::move(chPrimObj));
    }
    return currPrimObj;
}

}

inline UsdStage usdLoadStage(PythonClient &cli, std::string const &stageFile) {
    auto rootPrimPath = stageDefaultPrim(cli, stageFile);
    auto rootPrimObj = details::recursivelyLoadPrimAndItsChildren(cli, stageFile, rootPrimPath);
    return UsdStage{std::move(rootPrimObj)};
}

}
